import rospy
import math
import numpy as np

from webots_ros.srv import set_float
from webots_ros.srv import set_int
from webots_ros.srv import get_float
from webots_ros.msg import Float64Stamped

import math
TIMESTEP = 32

class ArmManager():
    def __init__(self):

        self.joint2value = 0
        webots_ns = 'robot'

        joint1_client = '{}/joint1/{}'
        joint2_client = '{}/joint2/{}'
        joint3_client = '{}/joint3/{}'
        joint4_client = '{}/joint4/{}'
        joint5_client = '{}/joint5/{}'

        
        # create proxy methods to set motor positions
        rospy.wait_for_service(joint1_client.format(webots_ns, 'set_position'))
        self.j1_set_position = rospy.ServiceProxy(
            joint1_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint2_client.format(webots_ns, 'set_position'))
        self.j2_set_position = rospy.ServiceProxy(
            joint2_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint3_client.format(webots_ns, 'set_position'))
        self.j3_set_position = rospy.ServiceProxy(
            joint3_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint4_client.format(webots_ns, 'set_position'))
        self.j4_set_position = rospy.ServiceProxy(
            joint4_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        rospy.wait_for_service(joint5_client.format(webots_ns, 'set_position'))
        self.j5_set_position = rospy.ServiceProxy(
            joint5_client.format(webots_ns, 'set_position'), set_float, persistent=True)

        # create proxy methods to set motor velocities
        rospy.wait_for_service(joint1_client.format(webots_ns, 'set_velocity'))
        self.j1_set_velocity = rospy.ServiceProxy(
            joint1_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint2_client.format(webots_ns, 'set_velocity'))
        self.j2_set_velocity = rospy.ServiceProxy(
            joint2_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint3_client.format(webots_ns, 'set_velocity'))
        self.j3_set_velocity = rospy.ServiceProxy(
            joint3_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint4_client.format(webots_ns, 'set_velocity'))
        self.j4_set_velocity = rospy.ServiceProxy(
            joint4_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service(joint5_client.format(webots_ns, 'set_velocity'))
        self.j5_set_velocity = rospy.ServiceProxy(
            joint5_client.format(webots_ns, 'set_velocity'), set_float, persistent=True)

        rospy.wait_for_service("/robot/joint1_sensor/enable")
        self.joint1_enable = rospy.ServiceProxy("/robot/joint1_sensor/enable", set_int, persistent=True)
        self.joint1_enable(1)
        rospy.wait_for_service("/robot/joint2_sensor/enable")
        self.joint2_enable = rospy.ServiceProxy("/robot/joint2_sensor/enable", set_int, persistent=True)
        self.joint2_enable(1)
        rospy.wait_for_service("/robot/joint4_sensor/enable")
        self.joint4_enable = rospy.ServiceProxy("/robot/joint4_sensor/enable", set_int, persistent=True)
        self.joint4_enable(1)

        # rospy.Subscriber('/robot/joint2_sensor/value', Float64Stamped, self.Joint2Callback)
        # rospy.spin()
    
    def Joint2Callback(self,msg):
        self.joint2value = msg.data 

    def set_velocity(self, joint_positions = [0,0,0,0,0,0], joint_speed = 0.05):
        while((abs(self.joint2value - joint_positions[1])>0.001)):
            j1v = joint_speed if joint_speed<(abs(0.6*joint_speed*math.sin(self.joint2value))) else abs(0.7*joint_speed*math.sin(self.joint2value))
            
            self.j1_set_velocity(j1v)
            # print(j1v)
            self.j2_set_velocity(joint_speed)
            self.j3_set_velocity(joint_speed)
            self.j4_set_velocity(joint_speed)
            self.j5_set_velocity(joint_speed)

            self.j1_set_position(joint_positions[0])
            self.j2_set_position(joint_positions[1])
            self.j3_set_position(joint_positions[2])
            self.j4_set_position(joint_positions[3])
            self.j5_set_position(joint_positions[4])     

    def set_velocity_for_each_joint(self, joint_speed):
        self.j1_set_velocity(joint_speed[0]*np.sign(joint_speed[0]))
        self.j2_set_velocity(joint_speed[1]*np.sign(joint_speed[1]))
        self.j4_set_velocity(joint_speed[2]*np.sign(joint_speed[2]))

        self.j1_set_position(100*np.sign(joint_speed[0]))
        self.j2_set_position(100*np.sign(joint_speed[1]))
        self.j4_set_position(100*np.sign(joint_speed[2]))

    def go_ahead(self, scale):
        pass
        # self.set_velocity(l_speed = -scale, r_speed = -scale)

        
if __name__ == "__main__":

    rospy.init_node('ArmManager', anonymous= False)

    cm = ArmManager()
    cm.set_velocity()
